import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense


class NM(object):
    def __init__(self):
        self.df = pd.read_csv('./timing/scenic_data.csv')

    def create_dataset(self, df, n_steps):
        """构建数据"""
        X, y = [], []
        for i in range(len(df) - n_steps):
            X.append(df.values[i:i + n_steps])
            y.append(df['count'].values[i + n_steps - 1])  # 假设这里'count'是目标变量，根据实际情况修改
        X = np.array(X)
        y = np.array(y)
        return X, y

    def get_model(self):
        """获取模型相关数据（这里只是简单示例，实际应用需完善）"""
        n_steps = 7  # 长度七天，可根据需求调整
        X, y = self.create_dataset(self.df, n_steps)

        # 训练集与测试集划分
        train_size = int(len(X) * 0.8)
        X_train, X_test = X[:train_size], X[train_size:]
        y_train, y_test = y[:train_size], y[train_size:]

        # 模型定义
        model = tf.keras.models.Sequential()
        model.add(tf.keras.layers.LSTM(50, activation='relu', return_sequences=True, input_shape=(n_steps, X.shape[2])))
        model.add(tf.keras.layers.LSTM(50, activation='relu'))
        model.add(tf.keras.layers.Dense(1))

        # 编译模型
        model.compile(optimizer='adam', loss='mse')

        # 训练模型
        model.fit(X_train, y_train, epochs=50, validation_data=(X_test, y_test))

        # 评估模型
        loss = model.evaluate(X_test, y_test)
        print(f"测试集损失: {loss}")

        return model